package io.kwdb.sdk.internal;

import io.kwdb.sdk.Config;
import java.util.*;
import java.util.concurrent.*;
import java.util.function.Consumer;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * Batches items for efficient bulk processing. Supports configurable batch size, linger time, and
 * backpressure strategies.
 *
 * @param <T> The type of items to batch
 */
public class Batcher<T> implements AutoCloseable {

  private static final Logger logger = Logger.getLogger(Batcher.class.getName());

  private final int maxBatchSize;
  private final long lingerMs;
  private final int maxQueueSize;
  private final Config.Backpressure backpressure;
  private final Consumer<List<T>> submitFunction;

  private final LinkedBlockingQueue<T> queue;
  private final List<T> currentBatch;
  private final Object batchLock = new Object();
  private final Thread workerThread;
  private final Metrics metrics;

  private volatile boolean closed = false;
  private volatile long lastBatchTime;

  /**
   * Create a new batcher with specified configuration.
   *
   * @param maxBatchSize Maximum items per batch (triggers submission)
   * @param lingerMs Maximum time to wait before submitting partial batch
   * @param maxQueueSize Maximum queue size for pending items
   * @param backpressure Backpressure strategy when queue is full
   * @param submitFunction Function to submit completed batches
   */
  public Batcher(
      int maxBatchSize,
      long lingerMs,
      int maxQueueSize,
      Config.Backpressure backpressure,
      Consumer<List<T>> submitFunction) {
    this.maxBatchSize = maxBatchSize;
    this.lingerMs = lingerMs;
    this.maxQueueSize = maxQueueSize;
    this.backpressure = backpressure != null ? backpressure : Config.Backpressure.BLOCK;
    this.submitFunction = submitFunction;

    this.queue = new LinkedBlockingQueue<>(maxQueueSize);
    this.currentBatch = new ArrayList<>(maxBatchSize);
    this.metrics = Metrics.getInstance();
    this.lastBatchTime = System.currentTimeMillis();

    // Start worker thread
    this.workerThread = new Thread(this::workerLoop, "Batcher-Worker");
    this.workerThread.setDaemon(true);
    this.workerThread.start();

    logger.info(
        String.format(
            "Batcher started: maxBatchSize=%d, lingerMs=%d, maxQueueSize=%d, backpressure=%s",
            maxBatchSize, lingerMs, maxQueueSize, this.backpressure));
  }

  /**
   * Offer an item to the batcher.
   *
   * @param item The item to batch
   * @param timeoutMillis Maximum time to wait if blocking (for BLOCK strategy)
   * @return true if accepted, false if rejected/timed out
   */
  public boolean offer(T item, long timeoutMillis) {
    if (closed) {
      return false;
    }

    metrics.updateQueueSize(queue.size());

    try {
      boolean accepted;

      if (backpressure == Config.Backpressure.BLOCK) {
        // Block until space available or timeout
        accepted = queue.offer(item, timeoutMillis, TimeUnit.MILLISECONDS);
        if (!accepted) {
          logger.fine("Queue full, timed out after " + timeoutMillis + "ms");
        }
      } else { // DROP_OLDEST
        accepted = queue.offer(item);
        if (!accepted && !queue.isEmpty()) {
          // Drop oldest item and retry
          queue.poll(); // Remove oldest
          metrics.incrementDropped(1);
          accepted = queue.offer(item);
          logger.fine("Dropped oldest item due to full queue");
        }
      }

      return accepted;

    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      return false;
    }
  }

  /** Worker loop that processes the queue and triggers batch submissions. */
  private void workerLoop() {
    logger.fine("Worker thread started");

    while (!closed || !queue.isEmpty()) {
      try {
        // Calculate timeout based on linger time
        long timeout = calculateTimeout();

        // Wait for item with timeout
        T item = queue.poll(timeout, TimeUnit.MILLISECONDS);

        if (item != null) {
          synchronized (batchLock) {
            currentBatch.add(item);

            // Check if batch is full
            if (currentBatch.size() >= maxBatchSize) {
              submitCurrentBatch();
            }
          }
        }

        // Check if linger time expired
        if (shouldSubmitByTime()) {
          synchronized (batchLock) {
            if (!currentBatch.isEmpty()) {
              submitCurrentBatch();
            }
          }
        }

      } catch (InterruptedException e) {
        // Expected during shutdown
        break;
      } catch (Exception e) {
        logger.log(Level.WARNING, "Error in batcher worker", e);
      }
    }

    // Submit any remaining items
    synchronized (batchLock) {
      if (!currentBatch.isEmpty()) {
        submitCurrentBatch();
      }
    }

    logger.fine("Worker thread stopped");
  }

  /** Calculate timeout for queue poll based on linger time. */
  private long calculateTimeout() {
    if (lingerMs <= 0) {
      return 1000; // Default 1 second if linger disabled
    }

    long elapsed = System.currentTimeMillis() - lastBatchTime;
    long remaining = lingerMs - elapsed;
    return Math.max(1, Math.min(remaining, 100)); // Between 1ms and 100ms
  }

  /** Check if batch should be submitted based on time. */
  private boolean shouldSubmitByTime() {
    if (lingerMs <= 0) {
      return false; // Linger disabled
    }

    long elapsed = System.currentTimeMillis() - lastBatchTime;
    return elapsed >= lingerMs;
  }

  /** Submit the current batch. Must be called while holding batchLock. */
  private void submitCurrentBatch() {
    if (currentBatch.isEmpty()) {
      return;
    }

    List<T> batchToSubmit = new ArrayList<>(currentBatch);
    currentBatch.clear();
    lastBatchTime = System.currentTimeMillis();

    // Update metrics
    metrics.observeBatchSize(batchToSubmit.size());

    // Submit batch (release lock first to avoid deadlock)
    try {
      logger.fine("Submitting batch of " + batchToSubmit.size() + " items");
      submitFunction.accept(batchToSubmit);
    } catch (Exception e) {
      logger.log(Level.SEVERE, "Error submitting batch", e);
    }
  }

  /** Flush any pending items immediately. */
  public void flush() {
    if (closed) {
      return;
    }

    synchronized (batchLock) {
      // Drain queue into current batch
      queue.drainTo(currentBatch);

      // Submit if not empty
      if (!currentBatch.isEmpty()) {
        submitCurrentBatch();
      }
    }
  }

  /** Close the batcher and flush remaining items. */
  @Override
  public void close() {
    if (closed) {
      return;
    }

    closed = true;
    logger.info("Closing batcher...");

    // Interrupt worker to speed up shutdown
    workerThread.interrupt();

    try {
      // Wait for worker to finish
      workerThread.join(5000);
      if (workerThread.isAlive()) {
        logger.warning("Worker thread did not stop in time");
      }
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      logger.warning("Interrupted while waiting for worker thread");
    }

    // Final flush (in case worker missed something)
    synchronized (batchLock) {
      queue.drainTo(currentBatch);
      if (!currentBatch.isEmpty()) {
        submitCurrentBatch();
      }
    }

    logger.info("Batcher closed");
  }
}
